{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "44f9d7bd",
   "metadata": {},
   "source": [
    "# 계층적 에이전트 팀\n",
    "\n",
    "이 튜토리얼에서는 **계층적 에이전트 팀**을 구성하는 방법을 살펴봅니다.\n",
    "\n",
    "단일 에이전트나 단일 수준의 감독자(supervisor)로는 대응하기 힘든 복잡한 작업을 **계층적 구조**를 통해 분할하고, 각각의 하위 수준 감독자(supervisor)가 해당 영역에 특화된 작업자(Worker) 에이전트를 관리하는 방식을 구현합니다.\n",
    "\n",
    "이러한 계층적 접근 방식은 작업자가 너무 많아질 경우나, 단일 작업자가 처리하기 힘든 복잡한 작업을 효율적으로 해결하는 데 도움이 됩니다.  \n",
    "\n",
    "본 예제는 [AutoGen 논문](https://arxiv.org/abs/2308.08155)의 아이디어를 LangGraph를 통해 구현한 사례로, 웹 연구와 문서 작성이라는 두 가지 하위 작업을 서로 다른 팀으로 구성하고, 상위 및 중간 수준의 감독자를 통해 전체 프로세스를 관리하는 방법을 제시합니다.\n",
    "\n",
    "![](./assets/langgraph-multi-agent-team-supervisor.png)\n",
    "\n",
    "---\n",
    "\n",
    "**왜 계층적 에이전트 팀인가?**\n",
    "\n",
    "이전 Supervisor 예제에서는 하나의 supervisor node가 여러 작업자 노드에게 작업을 할당하고 결과를 취합하는 과정을 살펴보았습니다. 이 방식은 간단한 경우에 효율적입니다. 그러나 다음과 같은 상황에서는 계층적 구조가 필요할 수 있습니다.\n",
    "\n",
    "- **작업 복잡성 증가**: 단일 supervisor로는 한 번에 처리할 수 없는 다양한 하위 영역의 전문 지식이 필요할 수 있습니다.\n",
    "- **작업자 수 증가**: 많은 수의 작업자를 관리할 때, 단일 supervisor가 모든 작업자에게 직접 명령을 내리면 관리 부담이 커집니다.\n",
    "\n",
    "이러한 상황에서 상위 수준의 supervisor는 하위 수준의 **sub-supervisor** 들에게 작업을 할당하고, 각 **sub-supervisor** 는 해당 작업을 전문화된 작업자 팀에 재할당하는 계층적 구조를 구성할 수 있습니다.\n",
    "\n",
    "---\n",
    "\n",
    "**이 튜토리얼에서 다룰 내용**\n",
    "\n",
    "1. **도구 생성**: 웹 연구(Web Research) 및 문서 작성(Documentation)을 위한 에이전트 도구 정의    \n",
    "2. **에이전트 팀 정의**: 연구 팀 및 문서 작성 팀을 계층적으로 정의하고 구성  \n",
    "3. **계층 추가**: 상위 수준 그래프와 중간 수준 감독자를 통해 전체 작업을 계층적으로 조정  \n",
    "4. **결합**: 모든 요소를 통합하여 최종적인 계층적 에이전트 팀 구축\n",
    "\n",
    "---\n",
    "\n",
    "**참고**\n",
    "\n",
    "- [AutoGen 논문: Enabling Next-Gen LLM Applications via Multi-Agent Conversation (Wu et al.)](https://arxiv.org/abs/2308.08155)\n",
    "- [LangGraph - Multi-Agent 개념](https://langchain-ai.github.io/langgraph/concepts/multi_agent/)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2da4c8ca",
   "metadata": {},
   "source": [
    "## 환경 설정"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5916579d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# API 키를 환경변수로 관리하기 위한 설정 파일\n",
    "from dotenv import load_dotenv\n",
    "\n",
    "# API 키 정보 로드\n",
    "load_dotenv()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "209cb6c2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# LangSmith 추적을 설정합니다. https://smith.langchain.com\n",
    "# !pip install -qU langchain-teddynote\n",
    "from langchain_teddynote import logging\n",
    "\n",
    "# 프로젝트 이름을 입력합니다.\n",
    "logging.langsmith(\"CH17-LangGraph-Use-Cases\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c152a0f5",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_teddynote.models import get_model_name, LLMs\n",
    "\n",
    "# 최신 모델 이름 가져오기\n",
    "MODEL_NAME = get_model_name(LLMs.GPT4o)\n",
    "print(MODEL_NAME)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b0a8cb37",
   "metadata": {},
   "source": [
    "## 도구 생성\n",
    "\n",
    "각 팀은 하나 이상의 에이전트로 구성되며, 각 에이전트는 하나 이상의 도구를 갖추게 됩니다. 아래에서는 다양한 팀에서 사용할 모든 도구를 정의합니다.\n",
    "\n",
    "먼저 연구 팀을 살펴보겠습니다.\n",
    "\n",
    "**ResearchTeam 도구**\n",
    "\n",
    "ResearchTeam은 웹에서 정보를 찾기 위해 검색 엔진과 URL 스크래퍼를 사용할 수 있습니다. ResearchTeam의 성능을 향상시키기 위해 추가 기능을 아래에 자유롭게 추가할 수 있습니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8ca895a7",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "from langchain_community.document_loaders import WebBaseLoader\n",
    "from langchain_teddynote.tools.tavily import TavilySearch\n",
    "from langchain_core.tools import tool\n",
    "\n",
    "# 검색 도구 정의(TavilySearch)\n",
    "tavily_tool = TavilySearch(max_results=5)\n",
    "\n",
    "\n",
    "# 웹 페이지에서 세부 정보를 스크래핑하기 위한 도구 정의\n",
    "@tool\n",
    "def scrape_webpages(urls: List[str]) -> str:\n",
    "    \"\"\"Use requests and bs4 to scrape the provided web pages for detailed information.\"\"\"\n",
    "    # 주어진 URL 목록을 사용하여 웹 페이지 로드\n",
    "    loader = WebBaseLoader(\n",
    "        web_path=urls,\n",
    "        header_template={\n",
    "            \"User-Agent\": \"Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/102.0.0.0 Safari/537.36\",\n",
    "        },\n",
    "    )\n",
    "    docs = loader.load()\n",
    "\n",
    "    # 로드된 문서의 제목과 내용을 포함한 문자열 생성\n",
    "    return \"\\n\\n\".join(\n",
    "        [\n",
    "            f'<Document name=\"{doc.metadata.get(\"title\", \"\")}\">\\n{doc.page_content}\\n</Document>'\n",
    "            for doc in docs\n",
    "        ]\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "403a88c4",
   "metadata": {},
   "source": [
    "**문서 작성 팀 도구**\n",
    "\n",
    "다음으로, 문서 작성 팀이 사용할 도구(파일 접근 도구)를 정의합니다. \n",
    "\n",
    "이 도구는 에이전트가 파일 시스템에 접근할 수 있도록 하며, 이는 안전하지 않을 수 있습니다. 따라서, 사용에 주의가 필요합니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "438d04b3",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "from typing import Dict, Optional, List\n",
    "from typing_extensions import Annotated\n",
    "\n",
    "\n",
    "# 임시 디렉토리 생성 및 작업 디렉토리 설정\n",
    "WORKING_DIRECTORY = Path(\"./tmp\")\n",
    "\n",
    "# tmp 폴더가 없으면 생성\n",
    "WORKING_DIRECTORY.mkdir(exist_ok=True)\n",
    "\n",
    "\n",
    "# 아웃라인 생성 및 파일로 저장\n",
    "@tool\n",
    "def create_outline(\n",
    "    points: Annotated[List[str], \"List of main points or sections.\"],\n",
    "    file_name: Annotated[str, \"File path to save the outline.\"],\n",
    ") -> Annotated[str, \"Path of the saved outline file.\"]:\n",
    "    \"\"\"Create and save an outline.\"\"\"\n",
    "    # 주어진 파일 이름으로 아웃라인을 저장\n",
    "    with (WORKING_DIRECTORY / file_name).open(\"w\") as file:\n",
    "        for i, point in enumerate(points):\n",
    "            file.write(f\"{i + 1}. {point}\\n\")\n",
    "    return f\"Outline saved to {file_name}\"\n",
    "\n",
    "\n",
    "# 문서 읽기\n",
    "@tool\n",
    "def read_document(\n",
    "    file_name: Annotated[str, \"File path to read the document.\"],\n",
    "    start: Annotated[Optional[int], \"The start line. Default is 0\"] = None,\n",
    "    end: Annotated[Optional[int], \"The end line. Default is None\"] = None,\n",
    ") -> str:\n",
    "    \"\"\"Read the specified document.\"\"\"\n",
    "    # 주어진 파일 이름으로 문서 읽기\n",
    "    with (WORKING_DIRECTORY / file_name).open(\"r\") as file:\n",
    "        lines = file.readlines()\n",
    "    # 시작 줄이 지정되지 않은 경우 기본값 설정\n",
    "    if start is not None:\n",
    "        start = 0\n",
    "    return \"\\n\".join(lines[start:end])\n",
    "\n",
    "\n",
    "# 문서 쓰기 및 저장\n",
    "@tool\n",
    "def write_document(\n",
    "    content: Annotated[str, \"Text content to be written into the document.\"],\n",
    "    file_name: Annotated[str, \"File path to save the document.\"],\n",
    ") -> Annotated[str, \"Path of the saved document file.\"]:\n",
    "    \"\"\"Create and save a text document.\"\"\"\n",
    "    # 주어진 파일 이름으로 문서 저장\n",
    "    with (WORKING_DIRECTORY / file_name).open(\"w\") as file:\n",
    "        file.write(content)\n",
    "    return f\"Document saved to {file_name}\"\n",
    "\n",
    "\n",
    "# 문서 편집\n",
    "@tool\n",
    "def edit_document(\n",
    "    file_name: Annotated[str, \"File path of the document to be edited.\"],\n",
    "    inserts: Annotated[\n",
    "        Dict[int, str],\n",
    "        \"Dictionary where key is the line number (1-indexed) and value is the text to be inserted at that line.\",\n",
    "    ],\n",
    ") -> Annotated[str, \"File path of the edited document.\"]:\n",
    "    \"\"\"Edit a document by inserting text at specific line numbers.\"\"\"\n",
    "\n",
    "    # 주어진 파일 이름으로 문서 읽기\n",
    "    with (WORKING_DIRECTORY / file_name).open(\"r\") as file:\n",
    "        lines = file.readlines()\n",
    "\n",
    "    # 삽입할 텍스트를 정렬하여 처리\n",
    "    sorted_inserts = sorted(inserts.items())\n",
    "\n",
    "    # 지정된 줄 번호에 텍스트 삽입\n",
    "    for line_number, text in sorted_inserts:\n",
    "        if 1 <= line_number <= len(lines) + 1:\n",
    "            lines.insert(line_number - 1, text + \"\\n\")\n",
    "        else:\n",
    "            return f\"Error: Line number {line_number} is out of range.\"\n",
    "\n",
    "    # 편집된 문서를 파일에 저장\n",
    "    with (WORKING_DIRECTORY / file_name).open(\"w\") as file:\n",
    "        file.writelines(lines)\n",
    "\n",
    "    return f\"Document edited and saved to {file_name}\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "248d565b",
   "metadata": {},
   "source": [
    "다음은 코드 실행 도구인 `PythonREPLTool` 을 정의합니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "4a5faf77",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_experimental.tools import PythonREPLTool\n",
    "\n",
    "# PythonREPL 도구\n",
    "python_repl_tool = PythonREPLTool()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "64e51557",
   "metadata": {},
   "source": [
    "## 다중 에이전트 생성을 위한 유틸리티 함수 정의\n",
    "\n",
    "다음은 작업을 간결하게 수행하기 위한 몇 가지 유틸리티 함수를 생성하는 방법입니다.\n",
    "\n",
    "이때 이전 튜토리얼에서 다룬 `functools.partial` 함수를 사용하여 에이전트 노드를 생성합니다.\n",
    "\n",
    "1. `worker agent` 생성.\n",
    "2. `sub-graph`의 `supervisor` 생성."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "15598e9e",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langgraph.graph import START, END\n",
    "from langchain_core.messages import HumanMessage\n",
    "from langchain_openai.chat_models import ChatOpenAI\n",
    "from langgraph.prebuilt import create_react_agent\n",
    "\n",
    "\n",
    "# 에이전트 팩토리 클래스\n",
    "class AgentFactory:\n",
    "    def __init__(self, model_name):\n",
    "        self.llm = ChatOpenAI(model=model_name, temperature=0)\n",
    "\n",
    "    def create_agent_node(self, agent, name: str):\n",
    "        # 노드 생성 함수\n",
    "        def agent_node(state):\n",
    "            result = agent.invoke(state)\n",
    "            return {\n",
    "                \"messages\": [\n",
    "                    HumanMessage(content=result[\"messages\"][-1].content, name=name)\n",
    "                ]\n",
    "            }\n",
    "\n",
    "        return agent_node\n",
    "\n",
    "\n",
    "# LLM 초기화\n",
    "llm = ChatOpenAI(model=MODEL_NAME, temperature=0)\n",
    "\n",
    "# Agent Factory 인스턴스 생성\n",
    "agent_factory = AgentFactory(MODEL_NAME)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f4011486",
   "metadata": {},
   "source": [
    "아래는 `AgentFactory` 를 사용하여 에이전트 노드를 생성하는 예시입니다.\n",
    "\n",
    "예시에서는 검색 에이전트를 생성하는 방법을 살펴보겠습니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "be2f5962",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 에이전트 정의\n",
    "search_agent = create_react_agent(llm, tools=[tavily_tool])\n",
    "# 에이전트 노드 생성\n",
    "search_node = agent_factory.create_agent_node(search_agent, name=\"Searcher\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "56d3a6dd",
   "metadata": {},
   "source": [
    "다음은 팀 감독자(Team Supervisor)를 생성하는 함수입니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "a146aa89",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder\n",
    "from langchain_openai import ChatOpenAI\n",
    "from pydantic import BaseModel\n",
    "from typing import Literal\n",
    "\n",
    "\n",
    "def create_team_supervisor(model_name, system_prompt, members) -> str:\n",
    "    # 다음 작업자 선택 옵션 목록 정의\n",
    "    options_for_next = [\"FINISH\"] + members\n",
    "\n",
    "    # 작업자 선택 응답 모델 정의: 다음 작업자를 선택하거나 작업 완료를 나타냄\n",
    "    class RouteResponse(BaseModel):\n",
    "        next: Literal[*options_for_next]\n",
    "\n",
    "    # ChatPromptTemplate 생성\n",
    "    prompt = ChatPromptTemplate.from_messages(\n",
    "        [\n",
    "            (\"system\", system_prompt),\n",
    "            MessagesPlaceholder(variable_name=\"messages\"),\n",
    "            (\n",
    "                \"system\",\n",
    "                \"Given the conversation above, who should act next? \"\n",
    "                \"Or should we FINISH? Select one of: {options}\",\n",
    "            ),\n",
    "        ]\n",
    "    ).partial(options=str(options_for_next))\n",
    "\n",
    "    # LLM 초기화\n",
    "    llm = ChatOpenAI(model=model_name, temperature=0)\n",
    "\n",
    "    # 프롬프트와 LLM을 결합하여 체인 구성\n",
    "    supervisor_chain = prompt | llm.with_structured_output(RouteResponse)\n",
    "\n",
    "    return supervisor_chain"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "45b717eb",
   "metadata": {},
   "source": [
    "## 에이전트 팀 정의\n",
    "\n",
    "연구 팀(Research Team)과 문서 작성 팀(Doc Writing Team)을 정의합니다.\n",
    "\n",
    "### 연구 팀(Research Team)\n",
    "\n",
    "연구 팀은 `search agent`와 `web scraping`을 담당하는 `research_agent`라는 두 개의 작업자 노드를 가집니다. 이들을 생성하고 팀 감독자도 설정해 보겠습니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "41639e16",
   "metadata": {},
   "outputs": [],
   "source": [
    "import operator\n",
    "from typing import List, TypedDict\n",
    "from typing_extensions import Annotated\n",
    "\n",
    "from langchain_core.messages import BaseMessage, HumanMessage\n",
    "from langchain_openai.chat_models import ChatOpenAI\n",
    "from langgraph.prebuilt import create_react_agent\n",
    "\n",
    "\n",
    "# 상태 정의\n",
    "class ResearchState(TypedDict):\n",
    "    messages: Annotated[List[BaseMessage], operator.add]  # 메시지\n",
    "    team_members: List[str]  # 멤버 에이전트 목록\n",
    "    next: str  # Supervisor 에이전트에게 다음 작업자를 선택하도록 지시\n",
    "\n",
    "\n",
    "# LLM 초기화\n",
    "llm = ChatOpenAI(model=MODEL_NAME, temperature=0)\n",
    "\n",
    "# 검색 노드 생성\n",
    "search_agent = create_react_agent(llm, tools=[tavily_tool])\n",
    "search_node = agent_factory.create_agent_node(search_agent, name=\"Searcher\")\n",
    "\n",
    "# 웹 스크래핑 노드 생성\n",
    "web_scraping_agent = create_react_agent(llm, tools=[scrape_webpages])\n",
    "web_scraping_node = agent_factory.create_agent_node(\n",
    "    web_scraping_agent, name=\"WebScraper\"\n",
    ")\n",
    "\n",
    "# Supervisor 에이전트 생성\n",
    "supervisor_agent = create_team_supervisor(\n",
    "    MODEL_NAME,\n",
    "    \"You are a supervisor tasked with managing a conversation between the\"\n",
    "    \" following workers: Search, WebScraper. Given the following user request,\"\n",
    "    \" respond with the worker to act next. Each worker will perform a\"\n",
    "    \" task and respond with their results and status. When finished,\"\n",
    "    \" respond with FINISH.\",\n",
    "    [\"Searcher\", \"WebScraper\"],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9f9a53e2",
   "metadata": {},
   "source": [
    "다음으로 라우팅할 노드를 선택하는 함수를 정의합니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "9b2630c2",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_next_node(x):\n",
    "    return x[\"next\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b033a91f",
   "metadata": {},
   "source": [
    "### Research Team 그래프 생성"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a09fa76a",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_teddynote.graphs import visualize_graph\n",
    "from langgraph.graph import StateGraph\n",
    "from langgraph.checkpoint.memory import MemorySaver\n",
    "\n",
    "# 그래프 생성\n",
    "web_research_graph = StateGraph(ResearchState)\n",
    "\n",
    "# 노드 추가\n",
    "web_research_graph.add_node(\"Searcher\", search_node)\n",
    "web_research_graph.add_node(\"WebScraper\", web_scraping_node)\n",
    "web_research_graph.add_node(\"Supervisor\", supervisor_agent)\n",
    "\n",
    "# 엣지 추가\n",
    "web_research_graph.add_edge(\"Searcher\", \"Supervisor\")\n",
    "web_research_graph.add_edge(\"WebScraper\", \"Supervisor\")\n",
    "\n",
    "# 조건부 엣지 정의: Supervisor 노드의 결정에 따라 다음 노드로 이동\n",
    "web_research_graph.add_conditional_edges(\n",
    "    \"Supervisor\",\n",
    "    get_next_node,\n",
    "    {\"Searcher\": \"Searcher\", \"WebScraper\": \"WebScraper\", \"FINISH\": END},\n",
    ")\n",
    "\n",
    "# 시작 노드 설정\n",
    "web_research_graph.set_entry_point(\"Supervisor\")\n",
    "\n",
    "# 그래프 컴파일\n",
    "web_research_app = web_research_graph.compile(checkpointer=MemorySaver())\n",
    "\n",
    "# 그래프 시각화\n",
    "visualize_graph(web_research_app, xray=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "21af9f6e",
   "metadata": {},
   "source": [
    "`web_research_app` 을 실행합니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "d4f9fc0b",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.runnables import RunnableConfig\n",
    "from langchain_teddynote.messages import random_uuid, invoke_graph\n",
    "\n",
    "\n",
    "def run_graph(app, message: str, recursive_limit: int = 50):\n",
    "    # config 설정(재귀 최대 횟수, thread_id)\n",
    "    config = RunnableConfig(\n",
    "        recursion_limit=recursive_limit, configurable={\"thread_id\": random_uuid()}\n",
    "    )\n",
    "\n",
    "    # 질문 입력\n",
    "    inputs = {\n",
    "        \"messages\": [HumanMessage(content=message)],\n",
    "    }\n",
    "\n",
    "    # 그래프 실행\n",
    "    invoke_graph(app, inputs, config)\n",
    "\n",
    "    return app.get_state(config).values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "acf3ea7f",
   "metadata": {},
   "outputs": [],
   "source": [
    "output = run_graph(\n",
    "    web_research_app,\n",
    "    \"https://finance.naver.com/news 의 주요 뉴스 정리해서 출력해줘. 출처(URL) 도 함께 출력해줘.\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "297ac021",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 최종 결과 출력\n",
    "print(output[\"messages\"][-1].content)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "94d75718",
   "metadata": {},
   "source": [
    "### 문서 작성 팀(Doc Writing Team)\n",
    "\n",
    "이번에는 문서 작성 팀을 생성합니다. 이때, 각 agent에게 서로 다른 file-writing 도구에 대한 접근 권한을 부여합니다. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "893e58da",
   "metadata": {},
   "outputs": [],
   "source": [
    "import operator\n",
    "from typing import List, TypedDict, Annotated\n",
    "from pathlib import Path\n",
    "\n",
    "\n",
    "# 임시 디렉토리 생성 및 작업 디렉토리 설정\n",
    "WORKING_DIRECTORY = Path(\"./tmp\")\n",
    "WORKING_DIRECTORY.mkdir(exist_ok=True)  # tmp 폴더가 없으면 생성\n",
    "\n",
    "\n",
    "# 상태 정의\n",
    "class DocWritingState(TypedDict):\n",
    "    messages: Annotated[List[BaseMessage], operator.add]\n",
    "    team_members: str\n",
    "    next: str\n",
    "    current_files: str  # 현재 작업중인 파일\n",
    "\n",
    "\n",
    "# 상태 전처리 노드: 각각의 에이전트가 현재 작업 디렉토리의 상태를 더 잘 인식할 수 있도록 함\n",
    "def preprocess(state):\n",
    "    # 작성된 파일 목록 초기화\n",
    "    written_files = []\n",
    "\n",
    "    try:\n",
    "        # 작업 디렉토리 내의 모든 파일을 검색하여 상대 경로로 변환\n",
    "        written_files = [\n",
    "            f.relative_to(WORKING_DIRECTORY) for f in WORKING_DIRECTORY.rglob(\"*\")\n",
    "        ]\n",
    "    except Exception:\n",
    "        pass\n",
    "    # 작성된 파일이 없으면 상태에 \"No files written.\" 추가\n",
    "    if not written_files:\n",
    "        return {**state, \"current_files\": \"No files written.\"}\n",
    "\n",
    "    # 작성된 파일 목록을 상태에 추가\n",
    "    return {\n",
    "        **state,\n",
    "        \"current_files\": \"\\nBelow are files your team has written to the directory:\\n\"\n",
    "        + \"\\n\".join([f\" - {f}\" for f in written_files]),\n",
    "    }\n",
    "\n",
    "\n",
    "# LLM 초기화\n",
    "llm = ChatOpenAI(model=MODEL_NAME)\n",
    "\n",
    "# 문서 작성 에이전트 생성\n",
    "doc_writer_agent = create_react_agent(\n",
    "    llm,\n",
    "    tools=[write_document, edit_document, read_document],\n",
    "    state_modifier=\"You are a arxiv researcher. Your mission is to write arxiv style paper on given topic/resources.\",\n",
    ")\n",
    "context_aware_doc_writer_agent = preprocess | doc_writer_agent\n",
    "doc_writing_node = agent_factory.create_agent_node(\n",
    "    context_aware_doc_writer_agent, name=\"DocWriter\"\n",
    ")\n",
    "\n",
    "# 노트 작성 노드\n",
    "note_taking_agent = create_react_agent(\n",
    "    llm,\n",
    "    tools=[create_outline, read_document],\n",
    "    state_modifier=\"You are an expert in creating outlines for research papers. Your mission is to create an outline for a given topic/resources or documents.\",\n",
    ")\n",
    "context_aware_note_taking_agent = preprocess | note_taking_agent\n",
    "note_taking_node = agent_factory.create_agent_node(\n",
    "    context_aware_note_taking_agent, name=\"NoteTaker\"\n",
    ")\n",
    "\n",
    "\n",
    "# 차트 생성 에이전트 생성\n",
    "chart_generating_agent = create_react_agent(\n",
    "    llm, tools=[read_document, python_repl_tool]\n",
    ")\n",
    "context_aware_chart_generating_agent = preprocess | chart_generating_agent\n",
    "chart_generating_node = agent_factory.create_agent_node(\n",
    "    context_aware_chart_generating_agent, name=\"ChartGenerator\"\n",
    ")\n",
    "\n",
    "# 문서 작성 팀 감독자 생성\n",
    "doc_writing_supervisor = create_team_supervisor(\n",
    "    MODEL_NAME,\n",
    "    \"You are a supervisor tasked with managing a conversation between the\"\n",
    "    \" following workers:  ['DocWriter', 'NoteTaker', 'ChartGenerator']. Given the following user request,\"\n",
    "    \" respond with the worker to act next. Each worker will perform a\"\n",
    "    \" task and respond with their results and status. When finished,\"\n",
    "    \" respond with FINISH.\",\n",
    "    [\"DocWriter\", \"NoteTaker\", \"ChartGenerator\"],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2310e4cc",
   "metadata": {},
   "source": [
    "### Doc Writing Team 그래프 생성"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "eba701f6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 그래프 생성\n",
    "authoring_graph = StateGraph(DocWritingState)\n",
    "\n",
    "# 노드 정의\n",
    "authoring_graph.add_node(\"DocWriter\", doc_writing_node)\n",
    "authoring_graph.add_node(\"NoteTaker\", note_taking_node)\n",
    "authoring_graph.add_node(\"ChartGenerator\", chart_generating_node)\n",
    "authoring_graph.add_node(\"Supervisor\", doc_writing_supervisor)\n",
    "\n",
    "# 엣지 정의\n",
    "authoring_graph.add_edge(\"DocWriter\", \"Supervisor\")\n",
    "authoring_graph.add_edge(\"NoteTaker\", \"Supervisor\")\n",
    "authoring_graph.add_edge(\"ChartGenerator\", \"Supervisor\")\n",
    "\n",
    "# 조건부 엣지 정의: Supervisor 노드의 결정에 따라 다음 노드로 이동\n",
    "authoring_graph.add_conditional_edges(\n",
    "    \"Supervisor\",\n",
    "    get_next_node,\n",
    "    {\n",
    "        \"DocWriter\": \"DocWriter\",\n",
    "        \"NoteTaker\": \"NoteTaker\",\n",
    "        \"ChartGenerator\": \"ChartGenerator\",\n",
    "        \"FINISH\": END,\n",
    "    },\n",
    ")\n",
    "\n",
    "# 시작 노드 설정\n",
    "authoring_graph.set_entry_point(\"Supervisor\")\n",
    "\n",
    "# 그래프 컴파일\n",
    "authoring_app = authoring_graph.compile(checkpointer=MemorySaver())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "07a8259a",
   "metadata": {},
   "source": [
    "그래프를 시각화 합니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3f82555f",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_teddynote.graphs import visualize_graph\n",
    "\n",
    "# 그래프 시각화\n",
    "visualize_graph(authoring_app, xray=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "868d9c65",
   "metadata": {},
   "source": [
    "그래프를 실행하고 결과를 확인합니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dc641fb1",
   "metadata": {},
   "outputs": [],
   "source": [
    "output = run_graph(\n",
    "    authoring_app,\n",
    "    \"Transformer 의 구조에 대해서 심층 파악해서 논문의 목차를 한글로 작성해줘. \"\n",
    "    \"그 다음 각각의 목차에 대해서 5문장 이상 작성해줘. \"\n",
    "    \"상세내용 작성시 만약 chart 가 필요하면 차트를 작성해줘. \"\n",
    "    \"최종 결과를 저장해줘. \",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8f6be24f",
   "metadata": {},
   "source": [
    "## Super-Graph 생성\n",
    "\n",
    "이 설계에서는 **상향식 계획 정책**을 적용하고 있습니다. 이미 두 개의 그래프를 생성했지만, 이들 간의 작업을 어떻게 라우팅할지 결정해야 합니다.\n",
    "\n",
    "이를 위해 **Super Graph**를 정의하여 이전 두 그래프를 조정하고, 이 상위 수준 상태가 서로 다른 그래프 간에 어떻게 공유되는지를 정의하는 연결 요소를 추가할 것입니다."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "60adddf3",
   "metadata": {},
   "source": [
    "먼저, 총 감독자 노드를 생성합니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "3dda4e3b",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.messages import BaseMessage\n",
    "from langchain_openai.chat_models import ChatOpenAI\n",
    "\n",
    "# 기본 LLM으로 ChatOpenAI 인스턴스 생성\n",
    "llm = ChatOpenAI(model=MODEL_NAME)\n",
    "\n",
    "# 팀 감독자 노드 생성\n",
    "supervisor_node = create_team_supervisor(\n",
    "    MODEL_NAME,\n",
    "    \"You are a supervisor tasked with managing a conversation between the\"\n",
    "    \" following teams: ['ResearchTeam', 'PaperWritingTeam']. Given the following user request,\"\n",
    "    \" respond with the worker to act next. Each worker will perform a\"\n",
    "    \" task and respond with their results and status. When finished,\"\n",
    "    \" respond with FINISH.\",\n",
    "    [\"ResearchTeam\", \"PaperWritingTeam\"],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ffb93745",
   "metadata": {},
   "source": [
    "다음은 Super-Graph의 상태와 노드를 정의 합니다.\n",
    "\n",
    "Super-Graph 는 단순하게 Task 를 라우팅 하는 역할이 주를 이룹니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "54d4f4d3",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import TypedDict, List, Annotated\n",
    "import operator\n",
    "\n",
    "\n",
    "# 상태 정의\n",
    "class State(TypedDict):\n",
    "    messages: Annotated[List[BaseMessage], operator.add]\n",
    "    # 라우팅 결정\n",
    "    next: str\n",
    "\n",
    "\n",
    "# 마지막 메시지 반환 노드\n",
    "def get_last_message(state: State) -> str:\n",
    "    last_message = state[\"messages\"][-1]\n",
    "    if isinstance(last_message, str):\n",
    "        return {\"messages\": [HumanMessage(content=last_message)]}\n",
    "    else:\n",
    "        return {\"messages\": [last_message.content]}\n",
    "\n",
    "\n",
    "# 응답 종합 노드\n",
    "def join_graph(response: dict):\n",
    "    # 마지막 메시지를 추출하여 메시지 목록으로 반환\n",
    "    return {\"messages\": [response[\"messages\"][-1]]}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "61c6add6",
   "metadata": {},
   "source": [
    "### Super-Graph 정의\n",
    "\n",
    "이제 2개의 팀을 연결하는 Super-Graph를 정의합니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "ddacbb75",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 그래프 정의\n",
    "super_graph = StateGraph(State)\n",
    "\n",
    "# 노드 정의\n",
    "super_graph.add_node(\"ResearchTeam\", get_last_message | web_research_app | join_graph)\n",
    "super_graph.add_node(\"PaperWritingTeam\", get_last_message | authoring_app | join_graph)\n",
    "super_graph.add_node(\"Supervisor\", supervisor_node)\n",
    "\n",
    "# 엣지 정의\n",
    "super_graph.add_edge(\"ResearchTeam\", \"Supervisor\")\n",
    "super_graph.add_edge(\"PaperWritingTeam\", \"Supervisor\")\n",
    "\n",
    "# 조건부 엣지 추가: Supervisor 의 결정에 따라 다음 노드로 이동\n",
    "super_graph.add_conditional_edges(\n",
    "    \"Supervisor\",\n",
    "    get_next_node,\n",
    "    {\n",
    "        \"PaperWritingTeam\": \"PaperWritingTeam\",\n",
    "        \"ResearchTeam\": \"ResearchTeam\",\n",
    "        \"FINISH\": END,\n",
    "    },\n",
    ")\n",
    "\n",
    "# Supervisor 노드를 시작 노드로 설정\n",
    "super_graph.set_entry_point(\"Supervisor\")\n",
    "\n",
    "# 그래프 컴파일\n",
    "super_graph = super_graph.compile(checkpointer=MemorySaver())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "73759297",
   "metadata": {},
   "source": [
    "그래프를 시각화 합니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f7e39e0f",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_teddynote.graphs import visualize_graph\n",
    "\n",
    "# 그래프 시각화\n",
    "visualize_graph(super_graph, xray=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8e5f5ef6",
   "metadata": {},
   "outputs": [],
   "source": [
    "output = run_graph(\n",
    "    super_graph,\n",
    "    \"\"\"주제: multi-agent 구조를 사용하여 복잡한 작업을 수행하는 방법\n",
    "\n",
    "상세 가이드라인:  \n",
    "- 주제에 대한 Arxiv 논문 형식의 리포트 생성\n",
    "- Outline 생성\n",
    "- 각각의 Outline 에 대해서 5문장 이상 작성\n",
    "- 상세내용 작성시 만약 chart 가 필요하면 차트 생성 및 추가\n",
    "- 한글로 리포트 작성\n",
    "- 출처는 APA 형식으로 작성\n",
    "- 최종 결과는 .md 파일로 저장\"\"\",\n",
    "    recursive_limit=150,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "58313e30",
   "metadata": {},
   "source": [
    "마크다운 형식으로 최종 결과물을 출력합니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ea2c963d",
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.display import Markdown\n",
    "\n",
    "if hasattr(output[\"messages\"][-1], \"content\"):\n",
    "    display(Markdown(output[\"messages\"][-1].content))\n",
    "else:\n",
    "    display(Markdown(output[\"messages\"][-1]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1279e7a7",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(output[\"messages\"][-1])"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "langchain-kr-lwwSZlnu-py3.11",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
